package midend.IR_Optimizer;

import jdk.jshell.spi.ExecutionControl;
import midend.Dom_Analysis;
import midend.IR_Instruction.*;
import midend.IR_Module;
import midend.IR_Type.IR_Float;
import midend.IR_Type.IR_Int32;
import midend.IR_Type.IR_Pointer;
import midend.IR_Type.IR_Type;
import midend.IR_Value.*;
import midend.Use;

import java.util.*;

public class New_Mem_To_Reg {
    private IR_Module module;
    private Instruction curInstr;
    private Stack<IR_Value> stack;
    private ArrayList<IR_BaseBlock> defBlockList;
    private ArrayList<IR_BaseBlock> useBlockList;
    private HashMap<Instruction, Instruction> defInstrList;
    private ArrayList<Instruction> useInstrList;
    public static HashMap<IR_FuncBlock,Dom_Analysis >dom_frontiersmap;

    private Map<IR_BaseBlock, Set<IR_BaseBlock>> dom_frontiers;
    private IR_Type i32 = IR_Int32.IR_Int32();
    private IR_Type floattype = IR_Float.IR_Float();

    public New_Mem_To_Reg(IR_Module module) {
        this.module = module;
    }
    public void run() {
        module.buildCFG();
        dom_frontiersmap = new HashMap<>();
        for (IR_FuncBlock func : module.getFuncBlocks()) {
            processFunction(func);
        }
    }

    private void processFunction(IR_FuncBlock func) {
        Dom_Analysis dom_analysis = new Dom_Analysis(func);
        dom_frontiers = dom_analysis.getDominanceFrontiers();
        dom_frontiersmap.put(func,dom_analysis);

        for (IR_BaseBlock block : func.get_bb_List()) {
            List<Instruction> instructions = new ArrayList<>(block.get_Instruction_List());
            for (Instruction instruction : instructions) {
                if (isValidAllocInstruction(instruction)) {
                    insertPhi(instruction);
                    rename(func.get_bb_List().get(0));
                }
            }
        }
    }

    private boolean isValidAllocInstruction(Instruction instruction) {
        if (!(instruction instanceof Alloc_Instrction)) return false;

        IR_Type baseType = ((IR_Pointer) instruction.get_op_index(0).get_Type()).get_Base_Type();
        return baseType == i32 || baseType == floattype;
    }

    public void insertPhi(Instruction instruction) {
        curInstr = instruction;
        stack = new Stack<>();
        this.useBlockList = new ArrayList<>();
        this.defBlockList = new ArrayList<>();
        this.useInstrList = new ArrayList<>();
        this.defInstrList = new HashMap<>();

        for (Use use : instruction.get_op_index(0).getUseList()) {
            assert use.getUser() instanceof Instruction;
            Instruction user = (Instruction) use.getUser();
            if (user instanceof Store_Instruction) {
                defInstrList.put(user, instruction);
                if (!defBlockList.contains(user.get_bb())) {
                    defBlockList.add(user.get_bb());
                }
            } else if (user instanceof Load_Instruction) {
                useInstrList.add(user);
                if (!defBlockList.contains(user.get_bb()))
                    useBlockList.add(user.get_bb());
            }
        }
        HashSet<IR_BaseBlock> blockneedPhi = new HashSet<>();//需要phi的block
        Stack<IR_BaseBlock> defValBlock = new Stack<>();//变量定义块集合
        for (IR_BaseBlock block : defBlockList) {
            defValBlock.push(block);
        }
        while (!defValBlock.isEmpty()) {
            IR_BaseBlock bb = defValBlock.pop();
            for (IR_BaseBlock domblock : dom_frontiers.get(bb)) {
                if (!blockneedPhi.contains(domblock)) {
                    IR_Value var = instruction.get_op_index(0);
                    IR_Type valueType = ((IR_Pointer)var.get_Type()).get_Base_Type();
                    IR_Value resRegister = new IR_VirtualReg_Value("phi", valueType); // 这里的phi_x只是变量名
                    Phi_Instruction phi = new Phi_Instruction(Collections.singletonList(resRegister), domblock); // 这一步工作只是插入phi指令，还没有对phi指令的数据流进行分析
                    domblock.get_Instruction_List().addFirst(phi);

                    useInstrList.add(phi);
                    defInstrList.put(phi, instruction);

                    blockneedPhi.add(domblock);
                    if (!defBlockList.contains(domblock)) {
                        defValBlock.push(domblock);
                    }
                }
            }
        }
    }

    public void rename(IR_BaseBlock entryBlock) {
        int count = 0;
        ArrayList<Instruction> deleteInstructions = new ArrayList<>();
        for (Instruction instruction : entryBlock.get_Instruction_List()) {
            if (useInstrList.contains(instruction) && instruction instanceof Load_Instruction) {
                IR_Value newValue = stack.peek();
                instruction.get_op_index(0).replaceUse(newValue);
                deleteInstructions.add(instruction);
            } else if (defInstrList.containsKey(instruction) && instruction instanceof Store_Instruction) {
                stack.push(instruction.get_op_index(1));
                count++;
                deleteInstructions.add(instruction);
            } else if (defInstrList.containsKey(instruction) && instruction instanceof Phi_Instruction) {
                stack.push(instruction.get_op_index(0));
                count++;
            } else if (instruction.equals(curInstr)) {
                deleteInstructions.add(instruction);
            }
        }

        for (IR_BaseBlock nextBlock : entryBlock.getNextBlock()) {
            Instruction firstInstruction = nextBlock.get_Instruction_List().getFirst();
            if (useInstrList.contains(firstInstruction) && firstInstruction instanceof Phi_Instruction phi) {
                IR_Value option = stack.empty() ? new IR_Int_Value(0,i32) : stack.peek();
                phi.set_ComingValue(option, entryBlock);
            }
        }

        for (Instruction instruction : deleteInstructions) {
            instruction.removewithvalue();
        }

        for (IR_BaseBlock block : entryBlock.getChildBlock()) {
            if (block.equals(entryBlock)) {
                continue;
            }
            rename(block);
        }
        for (int i = 0; i < count; i++) {
            stack.pop();
        }
    }
}
